from langchain.llms import OpenAI
from langchain.chains import ConversationChain
from sentence_transformers import SentenceTransformer
import chromadb


class QAHelper:
    """
    一个问答机器人的助手类，可以根据用户的输入，返回答案。
    具体功能：
    1、有用户问答上下文记忆的功能
    2、使用向量数据库实现RAG，提供专业的知识
    3、结合self_ask_with_search，提供更新准确和全面的回答
    """

    def __init__(self, data_path, db_path='chroma_db'):
        self.data_path = data_path
        self.db_path = db_path
        self.model = OpenAI(base_url="https://bd-gateway-agent.sdpsg.101.com/openai/v1/"
                            , api_key=os.getenv("OPENAI_API_KEY")
                            , temperature=0)
        self.embedder = SentenceTransformer('all-mpnet-base-v2')
        self.client = chromadb.PersistentClient(path=self.db_path)  # 初始化ChromaDB客户端

        # 尝试获取集合，如果不存在则创建
        self.collection = self.client.get_collection("documents")
        if not self.collection:
            self.collection = self.client.create_collection("documents")

    def load_data(self):
        # 假设数据在文本文件中，每行是一个文档
        with open(self.data_path, 'r', encoding='utf-8') as file:
            documents = [line.strip() for line in file if line.strip()]

        # 为每个文档生成向量并存储到ChromaDB
        for doc in documents:
            vector = self.embedder.encode([doc])
            self.collection.add(documents=[doc], vectors=[vector])

    def answer_question(self, question):
        # 对问题进行向量化
        question_vector = self.embedder.encode([question])
        # 在ChromaDB中检索最相关的文档
        results = self.collection.query(vectors=[question_vector], n_results=1)
        if results:
            retrieved_doc = results[0]['document']
            # 使用LangChain的ConversationChain生成答案
            conversation = ConversationChain(llm=self.model, input_messages=[retrieved_doc])
            return conversation.predict(input=question)
        return "Sorry, I could not find a relevant answer."

    def update_memory(self, question, answer):
        # 更新数据库
        question_vector = self.embedder.encode([question])
        self.collection.add(documents=[question], vectors=[question_vector], metadata={'answer': answer})


# 使用示例
if __name__ == '__main__':
    qa_helper = QAHelper(data_path="./KnowledgeDocuments")
    response = qa_helper.answer_question("请解释一下RAG的工作原理。")
    print(response)
